#-*- coding:utf-8 -*-
from prettytable import PrettyTable
from colorama import init, Fore
import platform
 
def isWindowsSystem():
    return 'Windows' in platform.system()
 
def isLinuxSystem():
    return 'Linux' in platform.system()
    
init()
class TrainsCollection:
    # todo 添加车票价格
    header = '车次 车站 时间 历时 一等 二等 软卧 硬卧 硬座 无座'.split()

    def __init__(self, available_trains, options):
        """查询到的火车班次集合

        :param available_trains: 一个列表, 包含可用的火车班次, 是返回数据中'data'对应的value
        :param options: 查询选项, 高铁, 动车等
        """
        self.available_trains = available_trains
        self.options = options

    def _get_duration(self, raw_train):
        duration = raw_train.get('lishi').replace(':', u'小时') + u'分'
        if duration.startswith('00'):
            return duration[4:]
        if duration.startswith('0'):
            return duration[1:]
        return duration

    @property
    def trains(self):
        for raw_train in self.available_trains:
            raw_train = raw_train['queryLeftNewDTO'] # Add v1.899 use newDTO
            train_no = raw_train['station_train_code']
            initial = train_no[0].lower()
            if not self.options or initial in self.options:
                train = [
                        train_no,
                        '\n'.join([Fore.GREEN + raw_train['from_station_name'] + Fore.RESET,
                            Fore.RED + raw_train['to_station_name'] + Fore.RESET]),
                        '\n'.join([Fore.GREEN + raw_train['start_time'] + Fore.RESET,
                            Fore.RED + raw_train['arrive_time'] + Fore.RESET]),
                        self._get_duration(raw_train),
                        raw_train['zy_num'], # 一等座
                        raw_train['ze_num'], # 二等座
                        raw_train['rw_num'], # 软卧
                        raw_train['yw_num'], # 硬卧
                        raw_train['yz_num'], # 硬座
                        raw_train['wz_num'], # 无座
                        ]
                yield train
                
    def pretty_print(self):
        pt = PrettyTable()
        pt._set_field_names(self.header)
        for train in self.trains:
            pt.add_row(train)
        if isWindowsSystem():
            print str(pt).decode('utf-8').encode('gbk')
        else:
            print pt
